import torch
import numpy as np
from torch.utils.data import Dataset
import random
import time
import os
import re
from PIL import Image,ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

class FeatureDataset(Dataset):
    def __init__(self, t_file, i_file):
        test_text = np.load(t_file)
        self.test_data_text = torch.from_numpy(test_text["data"]).float()
        test_img = np.load(i_file)
        self.test_data_img = torch.from_numpy(test_img["data"]).squeeze().float()
        self.test_labels = torch.from_numpy(test_text["label"]).long()

    def __len__(self):
        return self.test_data_text.shape[0]

    def __getitem__(self, item):
        return self.test_data_text[item], self.test_data_img[item], self.test_labels[item]



# 预处理
def text_preprocessing(text):
    """
    - 删除实体@符号(如。“@united”)
    — 纠正错误(如:'&amp;' '&')
    @参数 text (str):要处理的字符串
    @返回 text (Str):已处理的字符串
    """
    # 去除 '@name'
    text = re.sub(r'(@.*?)[\s]', ' ', text)

    #  替换'&amp;'成'&'
    text = re.sub(r'&amp;', '&', text)

    # 删除尾随空格
    text = re.sub(r'\s+', ' ', text).strip()

    return text


class FakeNewsDataset(Dataset):

    def __init__(self, df, root_dir, image_transform, tokenizer, MAX_LEN,div_data = 1):
        """
        参数:
            csv_file (string):包含文本和图像名称的csv文件的路径
            root_dir (string):目录
            transform(可选):图像变换
        """
        self.csv_data = df
        self.root_dir = root_dir
        self.image_transform = image_transform
        self.tokenizer_bert = tokenizer
        self.MAX_LEN = MAX_LEN
        self.div_data = div_data

    def __len__(self):
        #return int(self.csv_data.shape[0]/1000)
        return int(self.csv_data.shape[0]/self.div_data)
    
    def pre_processing_BERT(self, sent):

        # 创建空列表储存输出
        input_ids = []
        attention_mask = []
        
        encoded_sent = self.tokenizer_bert.encode_plus(
            text=text_preprocessing(sent),  # 预处理
            add_special_tokens=True,        # `[CLS]`&`[SEP]`
            max_length=self.MAX_LEN,        # 截断/填充的最大长度
            padding='max_length',           # 句子填充最大长度
            # return_tensors='pt',          # 返回tensor
            return_attention_mask=True,     # 返回attention mask
            truncation=True
            )
        
        input_ids = encoded_sent.get('input_ids')
        attention_mask = encoded_sent.get('attention_mask')
        
        # 转换tensor
        input_ids = torch.tensor(input_ids)
        attention_mask = torch.tensor(attention_mask)
        
        return input_ids, attention_mask
     
        
    def __getitem__(self, idx):

        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        img_name = os.path.join(self.root_dir, self.csv_data.iloc[idx]['#2 ImageID'])

        image = Image.open(img_name)
        if image.mode!="RGB":
            image = image.convert("RGB")

        image = self.image_transform(image)
        
        text = self.csv_data.iloc[idx]['#3 String']
        tensor_input_id, tensor_input_mask = self.pre_processing_BERT(text)

        label = self.csv_data.iloc[idx]['#1 Label']

        #label = int(label)
        
        label = torch.tensor(label)

        return [tensor_input_id, tensor_input_mask],image, label